import argparse
import os
import glob
import io
import msoffcrypto


def decrypt_excel_file(file_path, output_path, password):
    """解密单个Excel文件并保存为无密码版本"""
    try:
        with open(file_path, 'rb') as f:
            office_file = msoffcrypto.OfficeFile(f)

            if not office_file.is_encrypted():
                return False, "文件未加密"

            office_file.load_key(password=password)
            decrypted = io.BytesIO()
            office_file.decrypt(decrypted)

            with open(output_path, 'wb') as out_f:
                decrypted.seek(0)
                out_f.write(decrypted.read())

            return True, "解密成功"
    except Exception as e:
        return False, str(e)


def batch_remove_password(input_dir, output_dir, password):
    """批量处理目录中的Excel文件"""
    file_count = 0
    success_count = 0

    # 支持递归查找Excel文件
    for file_path in glob.glob(os.path.join(input_dir, '**', '*.*'), recursive=True):
        if not file_path.lower().endswith(('.xls', '.xlsx')):
            continue

        file_count += 1
        # 保持原始目录结构
        rel_path = os.path.relpath(file_path, input_dir)
        output_path = os.path.join(output_dir, rel_path)
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # 处理文件
        status, message = decrypt_excel_file(file_path, output_path, password)
        if status:
            print(f"✓ 成功处理: {rel_path}")
            success_count += 1
        else:
            print(f"✕ 处理失败: {rel_path} - {message}")

    print(f"\n处理完成：共尝试 {file_count} 个文件，成功 {success_count} 个")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='批量去除Excel文件密码')
    parser.add_argument('input_dir', help='包含加密Excel文件的目录路径')
    parser.add_argument('output_dir', help='输出目录路径')
    parser.add_argument('password', help='Excel文件的解密密码')

    args = parser.parse_args()

    print("正在处理，请稍候...\n")
    batch_remove_password(
        args.input_dir,
        args.output_dir,
        args.password
    )
